#coding:utf-8
from pymouse import PyMouse
from PIL import ImageGrab, Image
import string, random
#import cv2
import os
import numpy as np
from glob import glob
import math
from matplotlib import pyplot
import pickle

save_dir = 'F:\\lai\\projects\\xunyi\\path\\digit\\train_data_0'


lat_coarse = (1252, 1017, 1337, 1036)
log_coarse = (1415, 1017, 1507, 1036)
alt_coarse = (1581, 1017, 1623, 1036)

mouse_top_left = (343, 121)
mouse_down_right = (1720, 904)

m = PyMouse()

def crop_digits():
    for i in range(mouse_top_left[0], mouse_top_left[0]+1, 1):
        for j in range(mouse_top_left[1], mouse_top_left[1]+80, 3):
            m.move(i,j)
            screenshot = ImageGrab.grab()
            #screenshot = pyautogui.screenshot()
            lat_img = screenshot.crop(lat_coarse)
            log_img = screenshot.crop(log_coarse)
            alt_img = screenshot.crop(alt_coarse)
            
            # lat top left (48,4), right down (83,17)
            # log top left (55,3), right down (90,16)
            lat_digit = lat_img.crop((48,4,84,18))
            log_digit = log_img.crop((55,4,91,18))
            alt_digit = alt_img.crop((5,4,41,18))
            
            lat_name = ''.join([(string.ascii_letters)[x] for x in random.sample(range(0,52),8)]) + '.jpg'
            log_name = ''.join([(string.ascii_letters)[x] for x in random.sample(range(0,52),8)]) + '.jpg'
            alt_name = ''.join([(string.ascii_letters)[x] for x in random.sample(range(0,52),8)]) + '.jpg'
            #lat_digit.save(os.path.join(save_dir, lat_name))
            #log_digit.save(os.path.join(save_dir, log_name))
            alt_digit.save(os.path.join(save_dir, alt_name))

def load_labeled_data(src_dir, ratio):
    """
    load labeled digit data
      src_dir: The dir of labeld digits
      ratio: train ratio
    return
      (x_train, y_train, x_test, y_test)
    """
    img_files = glob(os.path.join(src_dir, '*.jpg'))
    n_train = int(math.floor(len(img_files)*ratio))
    img_train = random.sample(img_files, n_train)
    img_test = set(img_files) - set(img_train)
    
    x_train = np.zeros((n_train*4, 126), dtype=np.float32)
    y_train = np.zeros((n_train*4, ), dtype=np.float32)
    x_test = np.zeros((len(img_test)*4,126), dtype=np.float32)
    y_test = np.zeros((len(img_test)*4,), dtype=np.float32)
    k = 0
    for im in img_train:
        #img = cv2.imread(im)
        #img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = Image.open(im)
        img = img.convert('L')
        img = np.array(img)
        img_s = np.hsplit(img, 4)
        #pyplot.subplot(1,4,1)
        #pyplot.imshow(img_s[0])
        #pyplot.subplot(1,4,2)
        #pyplot.imshow(img_s[1])
        #pyplot.subplot(1,4,3)
        #pyplot.imshow(img_s[2])
        #pyplot.subplot(1,4,4)
        #pyplot.imshow(img_s[3])
        #pyplot.show()
        digits_str = os.path.split(im)[-1].split('.')[0]
        for i in range(4):
            x_train[k] = img_s[i].flatten()
            y_train[k] = int(digits_str[i])
            k = k + 1
    k = 0
    for im in img_test:
        #img = cv2.imread(im)
        #img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = Image.open(im)
        img = img.convert('L')
        img = np.array(img)        
        img_s = np.hsplit(img, 4)        
        digits_str = os.path.split(im)[-1].split('.')[0]
        for i in range(4):
            x_test[k] = img_s[i].flatten()
            y_test[k] = int(digits_str[i])
            k = k + 1
    return x_train, y_train, x_test, y_test

if __name__ == '__main__':
    # run load labeld data
    label_src_dir = 'F:\\lai\\projects\\xunyi\\path\\digit\\train_data_example'
    train_ratio = 0.8
    x_train, y_train, x_test, y_test =load_labeled_data(label_src_dir, train_ratio)
    pickle.dump((x_train, y_train, x_test, y_test), open('digit_train_test.pkl', 'wb'))

    #crop digits
    #crop_digits()

